import cv2
import os
import numpy as np
import argparse
import glob
import shutil

class Cam:
    '''
    摄像头类
    可以决定是否使用畸变矫正，并可以设置摄像头分辨率。
    相机参数文件用 <清晰度>.npz 命名，否则无法正确使用。
    Attributes:
        cap (cv2.VideoCapture): OpenCV视频捕捉对象。
        definitions (dict): 可用的分辨率定义。
        use_para (int): 是否使用畸变矫正，1为使用，0为不使用。
        defi (str): 摄像头分辨率定义，默认为"720p"。
        self.mtx (np.ndarray): 相机内参矩阵。
        self.dist (np.ndarray): 相机畸变系数。
    Methods:
        cam_attributes(): 获取摄像头的属性信息。
        get_frame(): 获取摄像头当前帧。
    '''
    def __init__(self, cam_index=1, *,  use_parser=True, use_para=1, defi="720p"):
        self.use_para = use_para
        self.defi = defi
        self.cap = cv2.VideoCapture(cam_index)
        if not self.cap.isOpened():
            raise Exception(f"无法打开摄像头 {cam_index}")
        self.definitions = {"480p": (640, 480), "720p": (1280, 720), "1080p": (1920, 1080)}
        if use_parser:
            # 设置可传入参数
            parser = argparse.ArgumentParser()
            parser.add_argument("-u", "--undisort", type=int, default=1, help="To decide whether use undisort or not") # 是否使用畸变矫正，1使用，0不使用。
            parser.add_argument("-d", "--definition", type=str, default="720p", help="To decide what definition the camera uses")
            args = parser.parse_args()
            self.use_para = args.undisort
            self.defi = args.definition
        # 设置摄像头分辨率 
        self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, self.definitions[self.defi][1])
        self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, self.definitions[self.defi][0])

        parameter_path = f"{self.defi}.npz"
        if self.use_para and os.path.exists(parameter_path):
            print("存在相机参数文件：", parameter_path) 
            parameter = np.load(parameter_path)
            self.mtx, self.dist = parameter.values()
            assert self.mtx.shape == (3, 3), "相机内参矩阵形状应为 (3, 3)"
            assert self.dist.shape == (1, 5), "相机畸变系数形状应为 (1, 5)"
        else:
            self.use_para = 0
    def get_cam_attributes(self):
        """
        获取摄像头的属性信息。
        Attributes:
            width (int): 视频帧的宽度。
            height (int): 视频帧的高度。
            fps (float): 视频帧率。
        """
        self.shape = (int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH)), 
                           int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
        self.fps = self.cap.get(cv2.CAP_PROP_FPS)
    def get_frame(self):
        """
        获取摄像头当前帧。
        Returns:
            frame (np.ndarray): 当前帧图像。
        """
        ret, frame = self.cap.read()
        if not ret:
            raise Exception("无法读取摄像头帧")
        if self.use_para:
            frame = cv2.undistort(frame, self.mtx, self.dist)
        return frame

class IMGprocess:
    '''
    这是一个用来进行图片预处理的类。
    Attributes:
        output_path(str): 存储图片路径
        imgs(list->str): 读取到的图片的路径
    Methods:
        low_value(img, para=0.4): 降低图片亮度
        add_noise(img): 添加噪点
        add_shelter(img, w=None, h=None): 添加遮挡，不指定w和h时，遮挡大小由程序自动选择。
        save_imgs(*funcs): 保存处理后的图片'  
    Usage:
        ipsr = IMGprocess(input_folder_path, output_folder_path)
        ipsr.save_imgs(ipsr.LowValue, ipsr.AddNoise, ipsr.AddShelter)
    注意：input_folder_path中只有.jpg或.png格式的图片可以被读取。
    '''
    def __init__(self, input_folder_path: str, output_folder_path: str, *, folder_recreate=True):
        self.output_path = output_folder_path
        imgs = glob.glob(f"{input_folder_path}/*.[jp][pn]g")
        assert len(imgs) > 0, "打开的文件夹不存在.jpg或.png格式的文件"
        self.imgs = imgs
        if folder_recreate:
            shutil.rmtree(output_folder_path, ignore_errors=True)
        os.makedirs(output_folder_path, exist_ok=True)
    
    def low_value(self, img, para=0.4):
        '''
        降低图片亮度
        :param img: 输入的图片
        :param para: 亮度降低的比例，默认值为0.4
        :return: 处理后的图片
        '''
        dst = cv2.convertScaleAbs(img, alpha=para, beta=0)  # alpha: 对比度，beta: 亮度
        return dst

    def add_noise(self, img):
        '''
        添加噪点
        :param img: 输入的图片
        :return: 处理后的图片
        '''
        noise = np.random.normal(2, 2, img.shape).astype(np.uint8)  # 正态分布噪声
        dst = cv2.add(img, noise)  # 添加噪声
        return dst
    
    def add_shelter(self, img, w=None, h=None):
        '''
        添加遮挡
        :param img: 输入的图片
        :param w: 遮挡区域的宽度，默认为图片宽度的1/8
        :param h: 遮挡区域的高度，默认为图片高度的1/8
        :return: 处理后的图片
        '''
        if w is None and h is None:
            w = h = img.shape[0] // 8
        else:
            w = w if w is not None else h
            h = h if h is not None else w
        (x, y) = np.random.randint(w, img.shape[1] - w*2), np.random.randint(h, img.shape[0] - h*2) # type: ignore
        mask = np.ones_like(img)*255
        mask[y:y+80, x:x+80] = 0  # 在指定区域添加遮挡
        dst = cv2.bitwise_and(img, mask)  # 应用遮挡
        return dst

    def save_imgs(self, *funcs):
        '''
        一键处理图片并保存
        :param funcs: 处理函数列表，函数的顺序决定了处理的顺序
        :return: None
        '''
        for i, img_path in enumerate(self.imgs):
            img = cv2.imread(img_path)
            img_name = os.path.splitext(os.path.basename(img_path))[0]
            for j, func in enumerate(funcs):
                processed_img = func(img)
                cv2.imwrite(f"{self.output_path}/{img_name}_{j+1}.jpg", processed_img)
            cv2.imwrite(f"{self.output_path}/{img_name}_0.jpg", img)